import os
import json
import random
import shutil
import cv2 as cv
from pathlib import Path

# name_list = ['C_JDX', 'C_DX', 'C_DXXJ', 'C_YC', 'C_CLS', 'C_CLSDLJXJ', 'C_CLSDXXJ', 'C_CLSDXXJ_KT','C_CLSDXXJ_KTLS',
#              'C_CLSZXMJS', 'C_CLSZXMJXJ', 'C_CLSZZDC', 'C_DLJX', 'C_DLJXJ', 'C_JCX', 'C_JCXDLJXJ', 'C_JCXDXXJ',
#              'C_JCXZXMJS', 'C_JCXZXMJXJ', 'C_JCXZXMJXJW']
name_list = ['C_DG','C_CLS','C_QX_QWSG','C_DLJX','C_QX_YZSG','C_CLSYS','C_DLJXYS','C_SG','C_BHT_DT']
def convert_json_to_yolo_obb(json_path, img_width, img_height):
    """将JSON标注转换为YOLOv8-OBB格式的字符串"""
    with open(json_path, 'r') as f:
        data = json.load(f)

    yolo_lines = []
    for shape in data['shapes']:
        # 获取类别ID（这里假设您的JSON中有class_id字段，如果没有需要修改）
        if shape['label'] in name_list:
            class_id = name_list.index(shape['label'])

            # 获取旋转矩形的四个顶点坐标
            points = shape['points']
            if len(points) != 4:
                continue  # 如果不是四边形则跳过

            normalized_points = []
            for x, y in points:
                x_norm = x / img_width
                y_norm = y / img_height
                normalized_points.extend([x_norm, y_norm])

            # 转换为YOLO-OBB格式：class_id x1 y1 x2 y2 x3 y3 x4 y4
            line = f"{class_id} " + " ".join(map(str, normalized_points))
            yolo_lines.append(line)
    return "\n".join(yolo_lines)


def split_dataset(img_dir, json_dir, output_dir, ratios=(0.9, 0.1, 0.0)):
    """划分数据集"""
    # 创建输出目录结构
    splits = ['train', 'val', 'test']
    for split in splits:
        os.makedirs(os.path.join(output_dir, split, 'images'), exist_ok=True)
        os.makedirs(os.path.join(output_dir, split, 'labels'), exist_ok=True)

    # 获取所有图像文件（不带扩展名）
    img_files = [os.path.splitext(f)[0] for f in os.listdir(img_dir)]
    random.shuffle(img_files)  # 随机打乱

    # 计算划分点
    total = len(img_files)
    train_end = int(total * ratios[0])
    val_end = train_end + int(total * ratios[1])

    # 划分数据集
    split_files = {
        'train': img_files[:train_end],
        'val': img_files[train_end:val_end],
        'test': img_files[val_end:]
    }

    # 处理每个划分集
    for split, files in split_files.items():
        print(f"Processing {split} set with {len(files)} samples...")

        for base_name in files:
        # 原始文件路径
            img_src = os.path.join(img_dir, base_name + '.jpg')  # 假设是jpg格式
            print("1111111111", img_src)
            img = cv.imread(img_src)
            height, width = img.shape[:2]
            json_src = os.path.join(json_dir, base_name + '.json')

            # 目标文件路径
            img_dst = os.path.join(output_dir, split, 'images', base_name + '.jpg')
            label_dst = os.path.join(output_dir, split, 'labels', base_name + '.txt')

            # 复制图像文件
            shutil.copy(img_src, img_dst)

            # 转换并保存标签文件
            if os.path.exists(json_src):
            # 这里需要知道图像尺寸来归一化坐标（如果需要）
            # 如果使用绝对坐标，可以传递任意尺寸，函数中不做归一化
                yolo_obb = convert_json_to_yolo_obb(json_src, width, height)  # 传入1,1表示使用绝对坐标
            with open(label_dst, 'w') as f:
                f.write(yolo_obb)

            print("Dataset splitting completed!")

if __name__ == '__main__':
    # 配置路径
    img_dir = '/home/liuhw/lhw_share/DSG/1/img'  # 原始图像目录
    json_dir = '/home/liuhw/lhw_share/DSG/1/json'  # 原始JSON标注目录
    output_dir = '/home/liuhw/lhw_share/DSG/1/txt'  # 输出目录
    split_dataset(img_dir, json_dir, output_dir)